import torch

# def rotate_around_neighbors(coordinates):
#     """
#     Rotate each coordinate by 90 degrees around its neighboring coordinate below it.
#     coordinates: Tensor of shape [B, H, W, 2]
#     """
#     B, H, W, _ = coordinates.shape
#     # Create a new tensor for the rotated coordinates
#     rotated = torch.zeros_like(coordinates)

#     # Handling interior points
#     for i in range(H-1):
#         for j in range(W):
#             # Center is the point directly below
#             center_x, center_y = coordinates[:, i+1, j, 0], coordinates[:, i+1, j, 1]
#             # Current point
#             point_x, point_y = coordinates[:, i, j, 0], coordinates[:, i, j, 1]
#             # Translate to origin based on center, rotate, translate back
#             new_x = center_x - (point_y - center_y)
#             new_y = center_y + (point_x - center_x)
#             rotated[:, i, j, 0] = new_x
#             rotated[:, i, j, 1] = new_y

#     # Handling last row (assuming rotation around themselves or no rotation)
#     rotated[:, -1, :, :] = coordinates[:, -1, :, :]  # No rotation or change as an example

#     return rotated


import torch

def rotate_around_center(coordinates):
    """
    Rotate each coordinate by 90 degrees around the center of all coordinates.
    coordinates: Tensor of shape [B, H, W, 2]
    """
    B, H, W, _ = coordinates.shape
    # Create a new tensor for the rotated coordinates
    rotated = torch.zeros_like(coordinates)

    # Calculate the center point
    center_x = coordinates[:, :, :, 0].float().mean()
    center_y = coordinates[:, :, :, 1].float().mean()

    # Rotate each point around the center
    for i in range(H):
        for j in range(W):
            # Current point
            point_x, point_y = coordinates[:, i, j, 0], coordinates[:, i, j, 1]
            # Translate to origin based on center, rotate, translate back
            new_x = center_x - (point_y - center_y)
            new_y = center_y + (point_x - center_x)
            rotated[:, i, j, 0] = new_x
            rotated[:, i, j, 1] = new_y

    return rotated

# Example usage
B, H, W = 1, 4, 4  # Batch size, height, and width
# Create an example tensor with shape [B, H, W, 2]
coordinates = torch.tensor([[[[1, 2], [3, 4], [5, 6], [7, 8]],
                             [[9, 10], [11, 12], [13, 14], [15, 16]],
                             [[17, 18], [19, 20], [21, 22], [23, 24]],
                             [[25, 26], [27, 28], [29, 30], [31, 32]]]])

# Rotate all coordinates 90 degrees around the center
rotated_coordinates = rotate_around_center(coordinates)
print("Original Coordinates:\n", coordinates)
print("Rotated Coordinates:\n", rotated_coordinates)

import matplotlib.pyplot as plt

# 旋转前后的坐标
original_coordinates = coordinates[0].view(-1, 2).numpy()
rotated_coordinates = rotated_coordinates[0].view(-1, 2).numpy()

# 创建一个新的图像
plt.figure()

# 绘制旋转前的点，使用红色表示
plt.scatter(original_coordinates[:, 0], original_coordinates[:, 1], color='r', label='Original')

# 绘制旋转后的点，使用蓝色表示
plt.scatter(rotated_coordinates[:, 0], rotated_coordinates[:, 1], color='b', label='Rotated')

# 添加图例
plt.legend()

# 显示图像
plt.show()


import matplotlib.pyplot as plt

# 旋转前后的坐标
original_coordinates = coordinates[0].view(-1, 2).numpy()
rotated_coordinates = rotated_coordinates[0].view(-1, 2).numpy()

# 创建一个新的图像
plt.figure()

# 绘制旋转前的点，使用红色表示
plt.scatter(original_coordinates[:, 0], original_coordinates[:, 1], color='r', label='Original')

# 绘制旋转后的点，使用蓝色表示
plt.scatter(rotated_coordinates[:, 0], rotated_coordinates[:, 1], color='b', label='Rotated')

# 绘制连接旋转前后的点的线，使用绿色表示
for i in range(original_coordinates.shape[0]):
    plt.plot([original_coordinates[i, 0], rotated_coordinates[i, 0]], 
             [original_coordinates[i, 1], rotated_coordinates[i, 1]], color='g')

# 添加图例
plt.legend()

# 显示图像
plt.show()